{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Numina Olympiad Problems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract AMC Data from numina\n",
    "from concurrent.futures import as_completed\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "from datasets import load_dataset\n",
    "\n",
    "ds = load_dataset(\"AI-MO/NuminaMath-CoT\")\n",
    "# Filter for amc_aime problems\n",
    "olympiad = ds['train'].filter(lambda x: x['source'] == 'olympiads')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Filter for non-proof problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter for non-proof problems.\n",
    "import concurrent.futures\n",
    "from deepscaler.utils import call_gemini_llm\n",
    "from deepscaler.system_prompts import FILTER_PROOF_PROMPT\n",
    "\n",
    "def process_entry(entry):\n",
    "    output_dict = {}\n",
    "    # 1) Get the problem text\n",
    "    problem_text = entry['problem']\n",
    "    solution_text = entry['solution']\n",
    "    # 2) Call Gemini LLM\n",
    "    output_str = call_gemini_llm(f'Problem: {problem_text} \\n\\n Solution: {solution_text}', system_prompt=FILTER_PROOF_PROMPT)\n",
    "    if not output_str:\n",
    "        print(\"Gemini not happy.\")\n",
    "        return None\n",
    "    if '[[2]]' in output_str:\n",
    "        output_dict['problem'] = entry['problem']\n",
    "        output_dict['solution'] = entry['solution']\n",
    "        return {\n",
    "            'problem': entry['problem'],\n",
    "            'solution': entry['solution'],\n",
    "        }\n",
    "    else:\n",
    "        if '[[3]]' in output_str:\n",
    "            print(output_str, entry['problem'])\n",
    "    return output_dict\n",
    "\n",
    "# Suppose `olympiad` is your list of dictionaries\n",
    "subset = olympiad  \n",
    "results = []\n",
    "\n",
    "with concurrent.futures.ProcessPoolExecutor(max_workers=48) as executor:\n",
    "    # 1) Submit all jobs to the executor\n",
    "    futures = [executor.submit(process_entry, entry) for entry in subset]\n",
    "\n",
    "# 2) Process them as they complete, using tqdm for a progress bar\n",
    "for future in tqdm(as_completed(futures), total=len(futures), desc=\"Processing entries\"):\n",
    "    # Get the result for each completed future\n",
    "    result = future.result()\n",
    "    if result:\n",
    "        results.append(result)\n",
    "\n",
    "\n",
    "# Save final list as json\n",
    "with open(\"olympiad_no_proof.json\", \"w\") as f:\n",
    "    json.dump(results, f, indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Filter for Unique Olympiad Problems (outside of Omni_Math)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from deepscaler.utils import RAG\n",
    "import json\n",
    "\n",
    "with open(\"olympiad_no_proof.json\", 'r', encoding='utf-8') as f:\n",
    "    olympiad_data = json.load(f)\n",
    "\n",
    "with open(\"../raw/train/omni_math.json\", 'r', encoding='utf-8') as f:\n",
    "    omni_data = json.load(f)\n",
    "\n",
    "rag_searcher = RAG(docs=[d[\"problem\"] for d in omni_data])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter for olympiad problems that are not in the omni dataset\n",
    "from tqdm import tqdm\n",
    "\n",
    "filter_olympiad = []\n",
    "num_problems = 0\n",
    "\n",
    "counter = 0\n",
    "# Wrap olympiad_data with tqdm, optionally adding a description and total\n",
    "for d in tqdm(olympiad_data, desc=\"Filtering olympiad data\", total=len(olympiad_data)):\n",
    "    search_result = rag_searcher.top_k(d[\"problem\"], k=1)[0]\n",
    "    score = search_result[\"score\"]\n",
    "    if score > 0.93:\n",
    "        num_problems += 1\n",
    "    else:\n",
    "        filter_olympiad.append(d)\n",
    "    counter += 1\n",
    "    if counter %1000 == 0:\n",
    "        print(counter)\n",
    "# Save final list as json\n",
    "with open(\"olympiad_unique.json\", \"w\") as f:\n",
    "    json.dump(filter_olympiad, f, indent=2)\n",
    "num_problems"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extract Answers from each Problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import concurrent.futures\n",
    "from concurrent.futures import as_completed\n",
    "import json\n",
    "# Process dataset to prpoduce answers\n",
    "with open(\"olympiad_no_proof.json\", 'r', encoding='utf-8') as f:\n",
    "    olympiad = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Have gemini add the solutions to the numina dataset.\n",
    "# Filter for non-proof problems.\n",
    "from deepscaler.utils import call_gemini_llm\n",
    "from deepscaler.system_prompts import EXTRACT_SOLUTION_PROMPT\n",
    "\n",
    "def get_answer(entry):\n",
    "    # 1) Get the problem text\n",
    "    problem_text = entry['problem']\n",
    "    solution_text = entry['solution']\n",
    "    # 2) Call Gemini LLM\n",
    "    output_list = call_gemini_llm(f'Problem: {problem_text} \\n----\\n Solution: {solution_text}', system_prompt=EXTRACT_SOLUTION_PROMPT, n=3)\n",
    "    output_list = [o for o in output_list if \n",
    "                   'error' not in o and 'Error' not in o and 'Solution not found' not in o]\n",
    "    if not output_list:\n",
    "        print(problem_text)\n",
    "        print(solution_text)\n",
    "        return None\n",
    "\n",
    "    for output_str in output_list:\n",
    "        if 'answer' in entry:\n",
    "            if output_str in entry['answer']:\n",
    "                continue\n",
    "            entry['answer'].append(output_str)\n",
    "        else:\n",
    "            entry['answer'] = [output_str]\n",
    "    return entry\n",
    "\n",
    "results = []\n",
    "idx = 0\n",
    "with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:\n",
    "    # 1) Submit all jobs to the executor\n",
    "    futures = [executor.submit(get_answer, entry) for entry in olympiad]\n",
    "    # 2) Process them as they complete, using tqdm for a progress bar\n",
    "    for future in as_completed(futures):\n",
    "        # Get the result for each completed future\n",
    "        result = future.result()\n",
    "        if result:\n",
    "            results.append(result)\n",
    "        if idx%1000 == 0:\n",
    "            print(idx)\n",
    "            with open(\"olympiad_checkpoint.json\", \"w\") as f:\n",
    "                json.dump(results, f, indent=2)\n",
    "        idx += 1\n",
    "\n",
    "# Save final list as json\n",
    "with open(\"olympiad_no_proof.json\", \"w\") as f:\n",
    "    json.dump(results, f, indent=2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "deepscaler",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
